package org.example.myleet.p480;

import java.util.*;

public class Solution {
    /**
     * 18 ms
     */
    public double[] medianSlidingWindow(int[] nums, int k) {
        double[] result = new double[nums.length-k+1];
        SlidingWindow window = new SlidingWindow(k);
        for (int i=0; i<k; i++) {
            //初始化，先塞满窗口
            window.insert(nums[i]);
        }
        for (int i=0; i<nums.length-k; i++) {
            //每次滑动获得中位数，并更新窗口内的元素
            result[i] = window.getMedian();
            window.insert(nums[i+k]);
            window.remove(nums[i]);
        }
        //最后滑动结束时填充最后一个结果
        result[result.length-1] = window.getMedian();
        return result;
    }

    static class SlidingWindow {
        int len;
        int smallerSize;
        int largerSize;
        PriorityQueue<Integer> smaller;
        PriorityQueue<Integer> larger;
        Map<Integer, Integer> delayDelete;

        public SlidingWindow(int l) {
            //窗口大小
            len = l;
            //较小部分和较大部分的有效size
            smallerSize = 0;
            largerSize = 0;
            //窗口中较小的部分，使用大根堆（倒序）
            smaller = new PriorityQueue<>(new Comparator<Integer>() {
                @Override
                public int compare(Integer o1, Integer o2) {
                    return o2.compareTo(o1);
                }
            });
            //窗口中较大的部分，使用小根堆（顺序）
            larger = new PriorityQueue<>(new Comparator<Integer>() {
                @Override
                public int compare(Integer o1, Integer o2) {
                    return o1.compareTo(o2);
                }
            });
            delayDelete = new HashMap<>();
        }

        public double getMedian() {
            //奇数时smaller的顶就是中位数，否则就是两边的和的均值，注意Integer.MAX相加溢出的情况，提前转成double避免溢出
            return (len & 1) == 1 ? smaller.peek() : (((double)smaller.peek() + (double)larger.peek()) / 2.0);
        }

        public void insert(int num) {
            //较小的放小的那边，较大的放大的那边
            if (smallerSize == 0 || num <= smaller.peek()) {
                smaller.offer(num);
                ++smallerSize;
            } else {
                larger.offer(num);
                ++largerSize;
            }
            balance();
        }

        public void remove(int num) {
            //记录需要删除的元素
            delayDelete.put(num, delayDelete.getOrDefault(num, 0)+1);
            if (num <= smaller.peek()) {
                --smallerSize;
                if (num == smaller.peek()) {
                    //可以被立即删除
                    prune(smaller);
                }
            } else {
                --largerSize;
                if (num == larger.peek()) {
                    //可以被立即删除
                    prune(larger);
                }
            }
            balance();
        }

        public void balance() {
            if (smallerSize > largerSize+1) {
                //小的部分比大的部分多了2个，从小的部分拿一个出来放在大的那边，最后两边都清理一下延迟删除的元素
                larger.offer(smaller.poll());
                --smallerSize;
                ++largerSize;
                prune(larger);
                prune(smaller);
            } else if (smallerSize < largerSize) {
                smaller.offer(larger.poll());
                --largerSize;
                ++smallerSize;
                prune(larger);
                prune(smaller);
            }
        }

        public void prune(PriorityQueue<Integer> heap) {
            //只要堆顶的元素属于延迟删除的，就进行删除
            Integer num = heap.peek();
            Integer count;
            while (delayDelete.containsKey(num)) {
                num = heap.poll();
                if ((count = delayDelete.get(num)) == 1) {
                    delayDelete.remove(num);
                } else {
                    delayDelete.put(num, count - 1);
                }
                num = heap.peek();
            }
        }
    }
}
